#include "img-utils.h"
#include <stdio.h>
#include <fstream>
#include <iostream>
#include <opencv2/opencv.hpp>

using namespace std;

#define CLIP0_255(temp) (temp < 0 ? 0 : (temp > 255 ? 255 : temp))

inline void yuv2pixel(uint8_t YUV[], uint8_t pixel[], bool rgb = true) {
    int R = 0, G = 1, B = 2;
    if (rgb == false) {
        B = 0;
        R = 2;
    }
    pixel[R] = CLIP0_255(YUV[0] + (((298799 * (YUV[2] - 128))) >> 18));
    pixel[G] = CLIP0_255(YUV[0] - ((103455 * (YUV[1] - 128) + 152201 * (YUV[2] - 128)) >> 18));
    pixel[B] = CLIP0_255(YUV[0] + (((532705 * (YUV[1] - 128))) >> 18));
}

inline void yuv2rgb(uint32_t height, uint32_t width, uint8_t* y, uint8_t* u, uint8_t* v, uint8_t* rgb) {
    uint8_t *up = u, *vp = v;
    uint8_t YUV[4];
    for (uint32_t h = 0; h < height; h += 2) {
        for (uint32_t w = 0; w < width; w += 2) {
            YUV[1] = *up++;
            YUV[2] = *vp++;
            YUV[0] = *(y + h * width + w);
            uint8_t* pixel = rgb + h * width * 3 + w * 3;
            yuv2pixel(YUV, pixel);
            YUV[0] = *(y + h * width + w + 1);
            yuv2pixel(YUV, pixel + 3);
            YUV[0] = *(y + (h + 1) * width + w);
            pixel = rgb + (h + 1) * width * 3 + w * 3;
            yuv2pixel(YUV, pixel);
            YUV[0] = *(y + (h + 1) * width + w + 1);
            yuv2pixel(YUV, pixel + 3);
        }
    }
}

inline void yuv2bgr(uint32_t height, uint32_t width, uint8_t* y, uint8_t* u, uint8_t* v, uint8_t* rgb) {
    uint8_t *up = u, *vp = v;
    uint8_t YUV[4];
    for (uint32_t h = 0; h < height; h += 2) {
        for (uint32_t w = 0; w < width; w += 2) {
            YUV[1] = *up++;
            YUV[2] = *vp++;
            YUV[0] = *(y + h * width + w);
            uint8_t* pixel = rgb + h * width * 3 + w * 3;
            yuv2pixel(YUV, pixel, false);
            YUV[0] = *(y + h * width + w + 1);
            yuv2pixel(YUV, pixel + 3, false);
            YUV[0] = *(y + (h + 1) * width + w);
            pixel = rgb + (h + 1) * width * 3 + w * 3;
            yuv2pixel(YUV, pixel, false);
            YUV[0] = *(y + (h + 1) * width + w + 1);
            yuv2pixel(YUV, pixel + 3, false);
        }
    }
}

void yuv420ColorCvt(unsigned char* src, int height, int width, unsigned char* dst, int code) {
    if (code < YUV2BGRA_NV21) {
        static cv::Mat yuvY(height, width, CV_8UC1);
        static cv::Mat yuvUV(height / 2, width / 2, CV_8UC2);
        static cv::Mat out(height, width, CV_8UC3);
        if (yuvY.rows != height || yuvY.cols != width) {
            yuvY = cv::Mat(height, width, CV_8UC1);
            yuvUV = cv::Mat(height / 2, width / 2, CV_8UC2);
            out = cv::Mat(height, width, CV_8UC3);
        }
        yuvY.data = src;
        yuvUV.data = src + width * height;
        out.data = dst;
        cv::cvtColorTwoPlane(yuvY, yuvUV, out, code);
    } else {
        unsigned char* y = src;
        unsigned char* u = src + height * width;
        unsigned char* v = u + height * width / 4;
        switch (code) {
            case YUV2RGB_I420: {
                yuv2rgb(height, width, y, u, v, dst);
                break;
            }
            case YUV2BGR_I420: {
                yuv2bgr(height, width, y, u, v, dst);
                break;
            }

            default:
                break;
        }
    }
}

void cvtBgr2Rgb(unsigned char* src, int height, int width) {
    unsigned char *p0, *p2;
    int size = height * width;
    int stride = 0;
    for (int i = 0; i < size; ++i) {
        p0 = src + stride;
        p2 = p0 + 2;
        *p0 = (*p0) ^ (*p2);
        *p2 = (*p0) ^ (*p2);
        *p0 = (*p0) ^ (*p2);
        stride += 3;
    }
}

void cvtBgr2Rgb(unsigned char* src, int height, int width, unsigned char* dst) {
    int size = height * width;
    int stride = 0;
    unsigned char xor_;
    for (int i = 0; i < size; ++i) {
        *(dst + stride + 1) = *(src + stride + 1);
        xor_ = *(src + stride) ^ *(src + stride + 2);
        *(dst + stride + 2) = xor_ ^ *(src + stride + 2);
        *(dst + stride) = xor_ ^ *(src + stride);
        stride += 3;
    }
}

void saveToFile(const char* path, unsigned char* data, unsigned int size) {
    std::ofstream outFile(path, std::ios::out | std::ios::binary);
    if (!outFile.is_open()) {
        std::cerr << "无法打开输出文件" << std::endl;
        throw std::bad_cast();
    }
    outFile.write(reinterpret_cast<const char*>(data), size);
    outFile.close();
}

void readFromFile(const char* path, unsigned char* buff, unsigned long size, int offset) {
    FILE* fp = fopen(path, "rb");
    if (NULL == fp) {
        std::cerr << "无法打开文件" << std::endl;
        throw std::bad_cast();
    }
    fseek(fp, offset, SEEK_SET);
    size_t ret = fread(buff, 1, size, fp);
    if (ret != size) {
        printf("WARN: read size is %ld\n", ret);
    }
    fclose(fp);
}